import logging
import numpy as np
from sklearn.datasets import fetch_openml
from sklearn.model_selection import train_test_split
from pinecone import pinecone
from tqdm import tqdm

logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(message)s')

# 加载MNIST数据集
mnist = fetch_openml('mnist_784', version=1)
X, y = mnist['data'], mnist['target']

# 划分训练集和测试集
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42)

# 初始化Pinecone客户端
pinecone.init(api_key="88add6a0-0d96-4779-8ecd-8d49b7da7226")

# 创建索引
index_name = "mnist-knn"
pinecone.create_index(index_name=index_name, metric="euclidean", shards=1)

# 获取索引实例
index = pinecone.Index(index_name=index_name)

# 上传训练数据到Pinecone
batch_size = 100
num_batches = len(X_train) // batch_size + 1
for i in tqdm(range(num_batches), desc="Uploading data to Pinecone"):
    start = i * batch_size
    end = min((i + 1) * batch_size, len(X_train))
    ids = [f"{idx}" for idx in range(start, end)]
    vectors = X_train[start:end].tolist()
    index.upsert(ids=ids, vectors=vectors)

logging.info("Successfully uploaded %d data points to Pinecone", len(X_train))

# 测试KNN准确率
def knn_accuracy(X_test, y_test, k):
    correct = 0
    for i in range(len(X_test)):
        results = index.query(queries=[X_test[i]], top_k=k)
        neighbors = [int(result['id']) for result in results['results'][0]['matches']]
        if y_test[i] in neighbors:
            correct += 1
    return correct / len(X_test)

k = 11
accuracy = knn_accuracy(X_test, y_test, k)
logging.info("Accuracy with k=%d using Pinecone: %.2f%%", k, accuracy * 100)
